package org.eoti.math;

import org.eoti.math.prime.PrimeUtil;
import org.eoti.util.StopWatch;

import java.math.BigInteger;
import java.util.List;
import java.util.concurrent.*;

/**
 * The MIT License
 * <p/>
 * Copyright (c) 2009 Malachi de AElfweald (malachid@gmail.com)
 * <p/>
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
 * documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to the following conditions:
 * <p/>
 * The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
 * Software.
 * <p/>
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
 * WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */
public class Shor
{
	private ScheduledExecutorService pool;
	private CopyOnWriteArrayList<ScheduledFuture<BigInteger>> tasks;
	private CopyOnWriteArrayList<BigInteger> factors;
	private BigInteger currentTotal = one;
	private BigInteger toFactor;
	private int threadNum = 0;

	public static void main(String[] args)
			throws ExecutionException, InterruptedException
	{
		if(args.length == 0)
		{
			System.out.format("USAGE: java org.eoti.math.Shor NUMBER\n");
			System.exit(1);
		}
		BigInteger toFactor = new BigInteger(args[0]);
		StopWatch watch = new StopWatch();
		watch.start();
		List<BigInteger> factors = factorize(toFactor);
		if(factors == null)
		{
			System.out.format("Unable to determine factors of %s in %s\n", toFactor, watch.report());
		}else{
			System.out.format("Factors of %s in %s\n", toFactor, watch.report());
			for(BigInteger factor : factors)
				System.out.format("\t%s\n", factor);
		}
	}

	public static List<BigInteger> factorize(BigInteger toFactor)
			throws ExecutionException, InterruptedException
	{
		Shor shor = new Shor(toFactor);
		if(shor.factors.size() == 0)
			return null;
		
		return shor.factors;
	}

	private Shor( BigInteger toFactor )
			throws ExecutionException, InterruptedException
	{
		this.toFactor = toFactor;
		tasks = new CopyOnWriteArrayList<ScheduledFuture<BigInteger>>();
		factors = new CopyOnWriteArrayList<BigInteger>();
		pool = Executors.newScheduledThreadPool( toFactor.bitLength() / 2 );

		// Simplify
		while(MathUtil.isEven( toFactor.divide(currentTotal) ))
			process(two);

		scheduleJobs();

		while(currentTotal == null || (!currentTotal.equals(toFactor)))
		{
			boolean found = false;

			for(ScheduledFuture<BigInteger> future : tasks)
			{
				if(future.isCancelled())
					continue;

				if(!future.isDone())
					continue;

				process(future.get());
				found = true;
			}

			if(found)
			{
				// This is in here just to cut down on how often it is executed
				if(currentTotal.equals(toFactor)) break;
				
				scheduleJobs();
			}
		}

		pool.shutdown(); // Disable new tasks from being submitted
		pool.shutdownNow(); // Cancel currently executing tasks		
	}

	protected boolean isPrime(BigInteger num)
	{
		if(!num.isProbablePrime(80)) return false;
		return PrimeUtil.isPrime(num);
	}

	protected void process(BigInteger factor)
	{
		BigInteger newTotal = currentTotal.multiply(factor);
		// This could happen if two threads got the same answer
		if(!toFactor.mod( newTotal ).equals( BigInteger.ZERO ))
			return;

//		System.out.format("factor: %s\n", factor);
		currentTotal = newTotal;
		factors.add(factor);
	}

	protected void scheduleJobs()
	{
		clearJobs();
		BigInteger subFactor = toFactor.divide(currentTotal);
		if(isPrime(subFactor))
			process(subFactor);
		else
			scheduleJobs(subFactor);
	}

	protected void scheduleJobs(BigInteger toSchedule)
	{
		int size = toSchedule.bitLength() / 2;
//		System.out.format("Scheduling %d jobs for %s\n", size, toSchedule);
		for(int i=0; i<size; i++)
			tasks.add(pool.schedule( new ShorTask(threadNum++, toSchedule), i, TimeUnit.MILLISECONDS ));
	}

	protected void clearJobs()
	{
		for(ScheduledFuture<BigInteger> future2 : tasks)
			future2.cancel( true );

		tasks.clear();
	}

	private static final BigInteger zero = BigInteger.ZERO;
	private static final BigInteger one = BigInteger.ONE;
	private static final BigInteger two = BigInteger.valueOf(2);

	static class ShorTask
	implements Callable<BigInteger>
	{
		private BigInteger N, a, g, r, b, c;
		private BigInteger toFactor, maxPeriod;
		private int threadNum;

		public ShorTask(int threadNum, BigInteger toFactor)
		{
//			System.out.format("ShorTask(%d,%s)\n", threadNum, toFactor);
			this.threadNum = threadNum;
			this.toFactor = toFactor;
			maxPeriod = BigInteger.valueOf(RandomUtil.randInt( 10 ) * 10 + 10);

			/**
			 * N: PQ
			 * a: random number < N
			 * g: gcd(a,N)
			 * r: period: smallest integer where f(x+r)=f(x)
			 * b: a^(r/2)
			 * c: (b+1)%N
			 * P: gcd(b+1, N)
			 * Q: gcd(b-1, N)
			 */
			N = toFactor;
			g = one;
			r = one;
			b = one;
			c = one;

		}

		/**
		 * Computes a result, or throws an exception if unable to do so.
		 *
		 * @return computed result
		 *
		 * @throws Exception if unable to compute a result
		 */
		public BigInteger call()
				throws Exception
		{
			while(true)
			{
				/**
				 * STEP 1: pick a random number a < N
				 */
				a = RandomUtil.randomBitFlip(toFactor);

				/**
				 * STEP 2: compute gcd(a,N)
				 */
				g = MathUtil.gcd(a, N);

				/**
				 * STEP 3: if g != 1, non-trivial factor, so we are done
				 */
				if(!g.equals(one))
					return g;

				/**
				 * STEP 4: compute r
				 */
				BigInteger exp = one;
				while(!exp.equals(maxPeriod))
				{
					// z: (a^exp)%N
					BigInteger z = a.modPow(exp, N);
					if(z.equals(one))
					{
						r = exp;
						break;
					}
					exp = exp.add(one);
				}

				/**
				 * STEP 5: if r is odd, go back to STEP 1
				 */
				if(r.mod(two).equals(one))
					continue;

				/**
				 * STEP 6: if a^(r/2) [congruent to (symbol won't display) -1 MOD N
				 */
				b = a.pow( (r.divide(two)).intValue() );
				c = (b.add(one)).mod(N);
				if(c.equals(zero))
					continue;

				/**
				 * STEP 7: factors
				 */
				BigInteger factor = MathUtil.gcd(b.add(one), N);
				if(factor.equals(one))
					factor = MathUtil.gcd(b.subtract(one), N);

				if(factor.equals(one))
					continue;

				return factor;
			}
		}
	}
}
